clc;
clear;

tic;
startTime = cputime;
pde = process_pde_2d(1);
vector_h = [1/4 1/8 1/16 1/32];
basis_type = "P1";
theta = 0;
disp(["basis_type: ", basis_type]);
if 0 == theta
    disp(["temporal_type: ", "forward Euler"]);
elseif 1 == theta
    disp(["temporal_type: ", "backward Euler"]);
elseif 0.5 == theta
    disp(["temporal_type: ", "Crank-Nicolson"]);
end
if pde.err
    error_Linf = zeros(size(vector_h,2),1);
    error_L2 = zeros(size(vector_h,2),1);
    error_H1 = zeros(size(vector_h,2),1);
end
% 由于cor函数和时间无关
% A不随时间变化
% 不区分A1, A2
for i = 1:size(vector_h,2)
    h = vector_h(i);
    mesh = process_mesh_2d(pde, h, basis_type);
    U1 = pde.u(mesh.node(:,1),mesh.node(:,2),0);
    M = assemble_mass_matrix_2d(mesh);
    A = assemble_stiffness_matrix_2d(mesh, pde, 0);
    if 1 == theta
        dt = 4*h^2;
    elseif 0.5 == theta
        dt = h;
    elseif 0 == theta
        error("stop");
    end
    Nt = (pde.domain(6) - pde.domain(5))/dt;
    for m = 0:Nt-1
        t1 = m * dt;
        t2 = (m+1) * dt;
        % A1 = assemble_stiffness_matrix_2d(mesh, pde, t1);
        % A2 = assemble_stiffness_matrix_2d(mesh, pde, t2);
        b1 = assemble_load_vector_2d(mesh, pde, t1);
        b2 = assemble_load_vector_2d(mesh, pde, t2);
        % [A1, b1] = process_boundary_t_2d(A1, b1, mesh, pde, t1);
        % [A2, b2] = process_boundary_t_2d(A2, b2, mesh, pde, t2);
        [A, b1] = process_boundary_t_2d(A, b1, mesh, pde, t1);
        [~, b2] = process_boundary_t_2d(A, b2, mesh, pde, t2);
        % A_hat = M/dt + theta*A2;
        % b_hat = theta*b2 + (1-theta)*b1 + (M/dt)*U1 - (1-theta)*A1*U1;
        % U2 = A_hat\b_hat;
        % U1 = U2;
        A_hat = M/(theta*dt) + A;
        b_hat = theta*b2 + (1-theta)*b1 + (M/(theta*dt))*U1;
        U_hat = A_hat\b_hat;
        U2 = (U_hat - U1)/theta + U1;
        U1 = U2;
    end
    if pde.err
        error_Linf(i,1) = process_error_t_2d(mesh, pde, U2, "Linf");
        error_L2(i,1) = process_error_t_2d(mesh, pde, U2, "L2");
        error_H1(i,1) = process_error_t_2d(mesh, pde, U2, "H1");
    end
    fprintf("完成\t%d/%d\n", i, size(vector_h,2));
end
if pde.err
    process_convergence_h_2d(error_Linf, vector_h);
    process_convergence_h_2d(error_L2, vector_h);
    process_convergence_h_2d(error_H1, vector_h);
end
elapsedTOCTime = toc;
elapsedCPUTime = cputime - startTime;
disp(["TOC time(s)", num2str(elapsedTOCTime)]);
disp(["CPU time(s)", num2str(elapsedCPUTime)]);


%% assemble_mass_matrix_2d
function matrix_M = assemble_mass_matrix_2d(mesh)
matrix_M = sparse(mesh.Nb,mesh.Nb);
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type);
    for i = 1:mesh.Nlb
        for j = 1:mesh.Nlb
            f = phi{j} .* phi{i};
            r = gauss_weights_loc * f;
            matrix_M(mesh.elem(n,j), mesh.elem(n,i)) = matrix_M(mesh.elem(n,j), mesh.elem(n,i)) + r;
        end
    end
end
end

%% assemble_stiffness_matrix_2d
function matrix_A = assemble_stiffness_matrix_2d(mesh, pde, t)
matrix_A = sparse(mesh.Nb,mesh.Nb);
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phix = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "dx");
    phiy = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "dy");
    fc = pde.cor(gauss_points_loc(:,1), gauss_points_loc(:,2), t);
    for i = 1:mesh.Nlb
        for j = 1:mesh.Nlb
            f = fc .* (phix{j} .* phix{i} + phiy{j} .* phiy{i});
            r = gauss_weights_loc * f;
            matrix_A(mesh.elem(n,j), mesh.elem(n,i)) = matrix_A(mesh.elem(n,j), mesh.elem(n,i)) + r;
        end
    end
end
end

%% assemble_load_vector_2d
function vector_b = assemble_load_vector_2d(mesh, pde, t)
vector_b = zeros(mesh.Nb,1);
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type);
    ff = pde.f(gauss_points_loc(:,1), gauss_points_loc(:,2), t);
    for j = 1:mesh.Nlb
        f =  ff .* phi{j};
        r = gauss_weights_loc * f;
        vector_b(mesh.elem(n,j),1) = vector_b(mesh.elem(n,j),1) + r;
    end
end
end